from torch.utils.data import Dataset
from torch import stack, as_tensor
from PIL import Image

class CustomDataset(Dataset):
    """
    花分类数据集
    """
    def __init__(self, dataset_path, dataset_labels, handle_methods = None):
        super().__init__()
        self.path = dataset_path
        self.labels = dataset_labels
        self.handle = handle_methods

    def __len__(self):
        return len(self.path)

    def __getitem__(self, item):
        """
        返回PIL图片数据与对应图片的label
        :param item: 图片索引
        :return: image, label
        """
        if self.handle is not None:
            return self.handle(Image.open(self.path[item])), self.labels[item]
        return Image.open(self.path[item]), self.labels[item]

    # 处理成pytorch识别的数据包
    @staticmethod
    def collate_func(batch):
        # -------------------------#
        #   batch数据格式如下
        #   |--batch
        #   |----item
        #   |------data
        #   |------label
        #   |----item
        #   ...
        #   ...
        #   |----item
        #   [(1, 2), (3, 4), (5, 6)]
        # -------------------------#

        # -------------------------------------------------------#
        #   将对象中对应的元素打包成一个个元组，然后返回由这些元组组成的列表
        #   [(1, 3, 5), (2, 4, 6)]
        #   也就是将图片数据和label一一对应
        # -------------------------------------------------------#
        images, labels = tuple(zip(*batch))
        # 将1x3x224x224的第一维进行堆叠
        images = stack(images, dim=0)
        # 可以接受任何像Python数据结构这样的数组
        labels = as_tensor(labels)
        return images, labels
